package m3f.data;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Random;

import m3f.computing.RandomNumbers;
import m3f.io.MatrixReader;
import m3f.io.MatrixReaderFactory;
import m3f.matrix.Minibatch;
import m3f.matrix.SparseVector;

public class MultimodalDataSet implements MultimodalDataProvider{
	
	private ArrayList<SparseVector> visual;
	private ArrayList<SparseVector> text;
	private int visualFeatures;
	private int textFeatures;
	private ArrayList<Integer> sample;
	private int visualSampleCounter;
	private int textSampleCounter;
	private int minibatchSize;
	
	private String visualFile;
	private String textFile;
	private boolean normalizeVectors;
	
	public MultimodalDataSet(String visualFile, String textFile, int visualFeatures, int textFeatures, boolean normalizeVectors){
		this.visualFeatures = visualFeatures;
		this.textFeatures = textFeatures;
		this.visualFile = visualFile;
		this.textFile  = textFile;
		this.normalizeVectors = normalizeVectors;
		init(visualFile, textFile, false, normalizeVectors);
	}
	
	public MultimodalDataSet(String visualFile, String textFile, boolean normalizeVectors){
		this.visualFile = visualFile;
		this.textFile  = textFile;
		this.normalizeVectors = normalizeVectors;	
		init(visualFile, textFile, true, normalizeVectors);	
	}
	
	public MultimodalDataSet clone(){
		MultimodalDataSet set = new MultimodalDataSet(visualFile, textFile, normalizeVectors);
		set.minibatchSize = this.minibatchSize;
		return set;
	}
	
	private void init(String visualFile, String textFile, boolean overrideFeatures, boolean normalizeVectors) {
		
		MatrixReader reader = MatrixReaderFactory.newInstance(visualFile, normalizeVectors);
		reader.start(visualFile);
		int visualRows = reader.getMatrixRows();
		if(overrideFeatures) {
			visualFeatures = reader.getMatrixColumns();
		}
		visual = reader.readVectors(visualRows);
		reader = MatrixReaderFactory.newInstance(textFile, normalizeVectors);
		reader.start(textFile);
		int textRows = reader.getMatrixRows();
		if(visualRows != textRows){
			System.err.println("ERROR: The data set is not aligned.");
			System.err.println("       Different number of vectors in visual and text sources.");
			System.err.println("       Visual: " + visual.size() + " Textual: " + text.size());
			System.exit(1);
		}
		if (overrideFeatures) {
			textFeatures = reader.getMatrixColumns();
		}
		text = reader.readVectors(textRows);
		reset(true);		
	}
	
	@Override
	public boolean hasMoreData() {
		return visualSampleCounter < visual.size();
	}

	public Minibatch maxTextBatch() {
		Minibatch tt = new Minibatch(textFeatures,text.size());
		for(int k = 0; k < text.size(); k++){
			tt.setColumn(k, text.get(k));
		}
		textSampleCounter = text.size();
		return tt;
	}
	
	public Minibatch maxVisualBatch() {
		Minibatch vv = new Minibatch(visualFeatures,visual.size());
		for(int k = 0; k < visual.size(); k++){
			vv.setColumn(k, visual.get(k));
		}
		visualSampleCounter = visual.size();
		return vv;	
	}

	@Override
	public Minibatch nextTextual() {
		int size = minibatchSize;
		if(textSampleCounter + minibatchSize > text.size()){
			size = text.size() - textSampleCounter;
		}
		Minibatch tt = new Minibatch(textFeatures, size);
		for(int k = 0; k < minibatchSize && textSampleCounter < text.size(); k++){
			tt.setColumn(k, text.get(sample.get(textSampleCounter)));
			textSampleCounter++;
		}
		return tt;
	}

	@Override
	public Minibatch nextVisual() {
		int size = minibatchSize;
		if(visualSampleCounter + minibatchSize > visual.size()){
			size = visual.size() - visualSampleCounter;
		}
		Minibatch vv = new Minibatch(visualFeatures, size);
		for(int k = 0; k < minibatchSize && visualSampleCounter < visual.size(); k++){
			vv.setColumn(k, visual.get(sample.get(visualSampleCounter)));
			visualSampleCounter++;
		}
		return vv;
	}

	@Override
	public void reset(boolean randomized){
		sample = new ArrayList<Integer>();
		for(int i = 0; i < visual.size(); i++){
			sample.add(i);
		}
		Random rnd = RandomNumbers.getInstance();
		if(randomized){
			Collections.shuffle(sample, rnd);
		}
		visualSampleCounter = 0;
		textSampleCounter = 0;
	}

	@Override
	public void setMinibatchSize(int size){
		minibatchSize = size;
	}
	
	@Override
	public int getMinibatchSize(){
		return minibatchSize;
	}
	
	public int matrixRows(){
		return visual.size();
	}

	@Override
	public int textFeatures() {
		return textFeatures;
	}
	
	@Override
	public int visualFeatures() {
		return visualFeatures;
	}
}
